# -*- coding:utf-8 -*-
# 远程控制，接受mocap位置，计算ik，发送ikq
from ctypes import *
import time
import numpy as np
import socket
import threading
from collections import deque
import struct
import mujoco

# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *



# 共享数据缓冲区
class SharedBuffer:
    def __init__(self):
        self.buffers = [deque(maxlen=1), deque(maxlen=1)]  # 双缓冲
        self.current = 0
        self.lock = threading.Lock()
    
    def write(self, data):
        with self.lock:
            self.buffers[1 - self.current].append(data)
    
    def read(self):
        with self.lock:
            self.current = 1 - self.current
        return self.buffers[self.current][-1] if self.buffers[self.current] else None

# 初始化缓冲区
listA_buffer = SharedBuffer()
listB_buffer = SharedBuffer()

# UDP配置
UDP_IP_A = "192.168.10.4"  # 需替换实际IP
UDP_PORT_SEND = 10001  # B的发送端口
UDP_PORT_RECV = 10000  # A的发送端口

# 发送线程（B->A）
def sender():
    sock_send = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    while True:
        data = listB_buffer.read()
        if data:
            packed = struct.pack('!I' + 'f'*len(data), len(data), *data)
            sock_send.sendto(packed, (UDP_IP_A, UDP_PORT_SEND))
        time.sleep(0.002)  # 500Hz

# 接收线程（A->B）
def receiver():
    sock_recv = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock_recv.bind(("0.0.0.0", UDP_PORT_RECV))
    while True:
        data, _ = sock_recv.recvfrom(1024)
        length = struct.unpack('!I', data[:4])[0]
        listA = struct.unpack('!' + 'f'*length, data[4:])
        listA_buffer.write(listA)




# ###########################################################################


# ###########################################################################


if __name__ == "__main__":


    
    # >>>>>初始化CAN卡部分>>>
    # 设置can or canfd接口
    lcanlib = LCAN()
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # open device
    ret = lcanlib.LCAN_OpenDevice(LCAN_USBCANFD2, devIndex)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANFD2 Open Device failed!")
        exit(0)
    # open success
    print("LYS USBCANFD2 Open Device Success!")


    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # 初始化canfd
    ret = lcanlib.LCAN_InitCANFD(LCAN_USBCANFD2, devIndex, canIndex, canfd_init_cfg)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("LYS USBCANFD2 Init Failed!")
    else:
        print("LYS USBCANFD2 Init Success!")
    

    # 清除can通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)
    print('sleeping')
    time.sleep(0.1)
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    motor_controllers = MotorController(motor_num)

    # 清除错误
    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_clear_err, motor_num)
    print('sleeping')
    time.sleep(0.1)

    # 上使能
    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_enable, motor_num)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("Enable Failed!")
    else:
        print("Enable Success!")
    

    # 在线更新标志位
    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_obj_init, motor_num)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("motor init Failed!")
    else:
        print("motor init Success!")



    # 设置p模式
    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_p_mode, motor_num)
    print('sleeping')
    time.sleep(0.1)
    if ret != LCAN_STATUS_OK:
        print("Set p mode failed!")
    else:
        print("Set p mode success!")


    
    # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    # 控制
    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)

    # 发送与读取>>>
    # P Mode

    pvc = np.zeros((motor_num,3))
    # 每个周期，发送的目标位置和当前角度误差，最大不超过max_delta_intq
    control_period = .005
    rad_per_sec = 100000 # 最大速度，每秒转角
    max_dq = int(80000) # 最大转角差

    max_delta_intq_eachloop = min(int(rad_per_sec*control_period),max_dq)

    
    time.sleep(3)
    # 控制总时间
    control_time = float(20)
    start = time.time()

    ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_obj_get, motor_num)
    q_real = receive_p_state(lcanlib, devIndex, canIndex)
    # 读取初始点位
    model = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_P.xml')) 
    data = mujoco.MjData(model)
    mocap_id = model.body("target").mocapid[0]
    mocap_pos = data.mocap_pos[mocap_id]
    mocap_quat = data.mocap_quat[mocap_id]

    # 启动服务
    threading.Thread(target=sender, daemon=True).start()
    threading.Thread(target=receiver, daemon=True).start()


    
    while time.time() - start < control_time:
        loop_start = time.time()

        q_real_float = trans_intq2q(q_real)

        

        
        # 读取listA（非阻塞）
        current_listA = listA_buffer.read()
        if current_listA:
            print(f"current_listA = {current_listA}")
            mocap_pos = np.array(current_listA[0:3])
            mocap_quat = np.array(current_listA[3:7])
        

        ikq_star = ik(q_real_float, mocap_pos, mocap_quat)

        # 往Adevice的发送数据
        listB_buffer.write(ikq_star.tolist())

        intikq = trans_q2intq(ikq_star) # 最终目标点（也可能是IK尚未收敛的中间变量）

        # 归零
        # intikq = np.zeros(motor_num)


        if np.max(np.abs(intikq-q_real)) > max_dq:
            intikq_clip = q_real + max_dq/np.max(np.abs(intikq - q_real)) * (intikq - q_real)

        else:
            intikq_clip = intikq

        # print(f"q_real = {q_real}")
        # print(f"intikq_clip = {intikq_clip}")
        print(f"intikq = {intikq}")

        ctrl_intq = get_interpoint(q_real, intikq_clip, max_delta_intq_eachloop)

        print(f"ctrl_intq = {ctrl_intq}")

        motor_controllers.set_p(ctrl_intq)
        # 发送指令
        # transfer_time1 = time.time()
        ret = lcanlib.LCAN_TransmitFD(LCAN_USBCANFD2, devIndex, canIndex, motor_controllers.can_set_p, motor_num)
        # print(f"transfer time cost: {time.time() - transfer_time1}")

        # 读取状态
        new_pvc = receive_pvc(lcanlib, devIndex, canIndex)

        if new_pvc is not None:
            pvc = new_pvc  # 仅当有有效数据时更新
            print(f"最新状态:\n {pvc}")
        else:
            print(f"get no pvc data\n")

        q_real = pvc.T[0]

        elapsed_time = time.time() - loop_start
        sleep_time = control_period - elapsed_time
        if sleep_time > 0:
            time.sleep(sleep_time)
            print(f"loop time cost: {elapsed_time}")
        else:
            print(f"loop overtime: {-sleep_time}")


    # 发送与读取<<<

    # 清除通信缓存区
    ret = lcanlib.LCAN_ClearBuffer(LCAN_USBCANFD2, devIndex, canIndex)
    # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    #Close Device
    lcanlib.LCAN_CloseDevice(LCAN_USBCANFD2, devIndex)
    print("LYS USBCANFD2 Close!")